#include "icmp.h"

#include <unistd.h> // for getpid系统调用
#include <netinet/ip_icmp.h> // for struct icmp, ICMP_ECHO
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <netinet/in.h>
#include <sys/time.h>
#include <iostream>
#include <arpa/inet.h>

// 发送 ICMP Echo Request
void send_icmp_echo_request(int sockfd, struct sockaddr_in *dest_addr, int seq) {
    struct icmp icmp_header;
    int packet_size = sizeof(struct icmp) + 56; // ICMP header + payload
    char packet[packet_size];

    // 初始化 ICMP 头
    memset(&icmp_header, 0, sizeof(icmp_header));
    icmp_header.icmp_type = ICMP_ECHO;
    icmp_header.icmp_code = 0;
    icmp_header.icmp_id = getpid() & 0xffff;
    icmp_header.icmp_seq = seq;
    memcpy(packet, &icmp_header, sizeof(icmp_header));

    // 填充一些数据到 payload
    memset(packet + sizeof(struct icmp), 0, 56); // 填充 56 字节的 payload

    // 计算 ICMP 校验和
    icmp_header.icmp_cksum = checksum((unsigned short *)packet, packet_size);
    memcpy(packet, &icmp_header, sizeof(icmp_header));

    // 发送 ICMP 包
    if (sendto(sockfd, packet, packet_size, 0, (struct sockaddr *)dest_addr, sizeof(*dest_addr)) < 0) {
        perror("sendto");
        exit(1);
    }
}

// 接收 ICMP Echo Reply
int receive_icmp_reply(int sockfd, struct sockaddr_in *dest_addr, int seq) {
    char buffer[1500];
    struct sockaddr_in recv_addr;
    socklen_t addr_len = sizeof(recv_addr);
    struct ip *ip_header;
    struct icmp *icmp_header;
    int bytes_received;
    struct timeval start_time, end_time;

    // 获取发送时间
    gettimeofday(&start_time, NULL);

    // 接收数据包
    while ((bytes_received = recvfrom(sockfd, buffer, sizeof(buffer), 0, (struct sockaddr *)&recv_addr, &addr_len)) > 0) {
        // 解析 IP 头
        ip_header = (struct ip *)buffer;
        icmp_header = (struct icmp *)(buffer + (ip_header->ip_hl << 2)); // 跳过 IP 头

        // 检查 ICMP 类型是否为 Echo Reply
        if (icmp_header->icmp_type == ICMP_ECHOREPLY && icmp_header->icmp_id == (getpid() & 0xffff) && icmp_header->icmp_seq == seq) {
            // 获取接收时间
            gettimeofday(&end_time, NULL);

            // 计算 RTT（Round-Trip Time）
            double rtt = (end_time.tv_sec - start_time.tv_sec) * 1000.0 + (end_time.tv_usec - start_time.tv_usec) / 1000.0;
            printf("接收到来自%s的ICMP应答报文: seq=%d, ttl=%d, rtt=%.2f ms\n",
                   inet_ntoa(dest_addr->sin_addr), icmp_header->icmp_seq, ip_header->ip_ttl, rtt);
            return 0; // 成功接收
        }
    }

    if (bytes_received < 0) {
        perror("recvfrom");
    }
    return -1; // 未收到响应
}

// 计算 ICMP 校验和
unsigned short checksum(unsigned short *buf, int len) {
    unsigned long sum = 0;
    while (len > 1) {
        sum += *buf++;
        len -= 2;
    }
    if (len == 1) {
        sum += *(unsigned char *)buf;
    }
    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    return (unsigned short)(~sum);
}

void icmpEcho(char* ip){

    int sockfd;
    struct sockaddr_in dest_addr;

    // 创建 Raw Socket
    if ((sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)) < 0) {
        perror("socket");
        exit(1);
    }

    // 设置目标地址
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_addr.s_addr = inet_addr(ip); // 目标 IP

    // 发送 ICMP Echo Request
    int seq = 1;
    send_icmp_echo_request(sockfd, &dest_addr, seq);

    // 接收 ICMP Echo Reply
    if (receive_icmp_reply(sockfd, &dest_addr, seq) != 0) {    
        std::cout << "host is down";
    } 
    close(sockfd);
}